
#ifndef __TFTP_PACKET_HPP__
#define __TFTP_PACKET_HPP__

////////////////////////////////////////////////////////////////////////////////

#include "sock-inet.hpp"
#include "sock-inet-udp.hpp"
#include "tftp-exception.hpp"

#include <fstream>
#include <stdint.h>
#include <stdlib.h>

////////////////////////////////////////////////////////////////////////////////

class tftp_packet
{
public:
    tftp_packet(const uint8_t* packet, size_t len);
    const uint8_t* get_packet() const;
    size_t get_len() const;
    uint16_t get_opcode() const;

    enum opcode { read = 1, write = 2, data = 3, ack = 4, error = 5 };
    enum defaults { size = 516, data_size = 512 };

protected:
    const uint8_t* packet;
    size_t len;
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_flow : public tftp_packet
{
protected:
    tftp_packet_flow(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg);
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_request : public tftp_packet_flow
{
protected:
    tftp_packet_request(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg);

public:
    const char* get_filename() const;
    const char* get_mode() const;

protected:
    const char* filename;
    const char* mode;
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_read : public tftp_packet_request
{
public:
    tftp_packet_read(const uint8_t* packet, size_t len);
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_write : public tftp_packet_request
{
public:
    tftp_packet_write(const uint8_t* packet, size_t len);
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_transmission : public tftp_packet_flow
{
protected:
    tftp_packet_transmission(const uint8_t* packet, size_t len, uint16_t expected_opcode, const char* invalid_opcode_msg, uint16_t expected_sn);

public:
    uint16_t get_sn() const;
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_ack : public tftp_packet_transmission
{
public:
    tftp_packet_ack(const uint8_t* packet, size_t len, uint16_t expected_sn);
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_data : public tftp_packet_transmission
{
public:
    tftp_packet_data(const uint8_t* packet, size_t len, uint16_t expected_sn);

    const uint8_t* get_data() const;
    size_t get_datalen() const;
};

std::ofstream& operator<<(std::ofstream& stream, const tftp_packet_data& p);

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_error : public tftp_packet_flow
{
public:
    tftp_packet_error(const uint8_t* packet, size_t len);
    uint16_t get_code() const;
    const char* get_msg() const;

protected:
    const char* msg;
};

////////////////////////////////////////////////////////////////////////////////

class tftp_packet_factory
{
public:
    tftp_packet_factory();
    tftp_packet_factory(const sock_inet::address& addr);
    const uint8_t* get_packet() const;
    uint8_t* get_packet();
    size_t get_len() const;
    const sock_inet::address& get_address() const;
    tftp_packet_factory& operator<<(uint16_t x);
    tftp_packet_factory& operator<<(const char* x);
    tftp_packet_factory& operator<<(std::ifstream& stream);

protected:
    uint8_t buffer[tftp_packet::size];
    size_t cursor;
    sock_inet::address addr;
};

sock_inet_udp& operator<<(sock_inet_udp& s, const tftp_packet_factory& p);

////////////////////////////////////////////////////////////////////////////////

#endif
